from datetime import datetime
import backtrader as bt
import matplotlib.pyplot as plt
import akshare as ak
import pandas as pd
from collections import deque

plt.rcParams["font.sans-serif"] = ["SimHei"]
plt.rcParams["axes.unicode_minus"] = False

# 获取股票后复权数据
stock_hfq_df = ak.stock_zh_a_hist(symbol="000002", adjust="hfq").iloc[:, :6]
# 处理字段命名
stock_hfq_df.columns = ['date', 'open', 'close', 'high', 'low', 'volume']
# 将date设为索引
stock_hfq_df.index = pd.to_datetime(stock_hfq_df['date'])

class MACDBacktest(bt.Strategy):
    params = (
        ("macd_fast", 12),  # MACD快线周期
        ("macd_slow", 26),  # MACD慢线周期
        ("macd_signal", 9),  # MACD信号线周期
        ("lookback", 5),  # 背离检查的周期数
    )

    def __init__(self):
        self.data_close = self.datas[0].close
        self.order = None
        self.buy_price = None
        self.buy_comm = None
        
        # 初始化MACD指标
        self.macd = bt.indicators.MACDHisto(self.datas[0],
                                           period_me1=self.params.macd_fast,
                                           period_me2=self.params.macd_slow,
                                           period_signal=self.params.macd_signal)
        
        # 用于记录价格和MACD值的deque，用于背离检测
        self.price_deque = deque(maxlen=self.params.lookback)
        self.macd_deque = deque(maxlen=self.params.lookback)

    def log(self, txt, dt=None):
        ''' Logging function for this strategy'''
        dt = dt or self.datas[0].datetime.date(0)
        print(f'{dt.isoformat()}, {txt}')

    def notify_order(self, order):
        if order.status in [order.Submitted, order.Accepted]:
            return  # 订单还在处理中，不采取行动

        if order.status in [order.Completed]:
            if order.isbuy():
                self.log(f'BUY EXECUTED --- Price: {order.executed.price:.2f}, Cost: {order.executed.value:.2f}, Commission: {order.executed.comm:.2f}')
                self.buy_price = order.executed.price
                self.buy_comm = order.executed.comm
            else:
                self.log(f'SELL EXECUTED --- Price: {order.executed.price:.2f}, Cost: {order.executed.value:.2f}, Commission: {order.executed.comm:.2f}')
            self.bar_executed = len(self)
            
        elif order.status in [order.Canceled, order.Margin, order.Rejected]:
            self.log('Order Canceled/Margin/Rejected')
            
        self.order = None

    def next(self):
        if self.order:
            return
        
        # 更新价格和MACD值的deque
        self.price_deque.append(self.data_close[0])
        macd_val = self.macd.lines.macd[0]  # 使用MACD线（快线-慢线）
        self.macd_deque.append(macd_val)
        
        # 检查底背离
        if self.data_close[0] < min(self.price_deque) and macd_val > min(self.macd_deque):
            self.log("底背离，买入")
            self.order = self.buy()
        
        # 检查顶背离
        elif self.data_close[0] > max(self.price_deque) and macd_val < max(self.macd_deque):
            self.log("顶背离，卖出")
            self.order = self.sell()

# 初始化回测系统
cerebro = bt.Cerebro()

# 设置回测时间范围
start_date = datetime(1991, 4, 3)
end_date = datetime(2024, 5, 22)

# 加载数据到Cerebro
data = bt.feeds.PandasData(dataname=stock_hfq_df, fromdate=start_date, todate=end_date)
cerebro.adddata(data)

# 加载策略到Cerebro
cerebro.addstrategy(MACDBacktest)

# 设置初始资金和手续费
start_cash = 1000000
cerebro.broker.setcash(start_cash)
cerebro.broker.setcommission(commission=0.002)

# 运行回测
cerebro.run()

# 获取回测结果并打印
port_value = cerebro.broker.getvalue()
pnl = port_value - start_cash
print(f"初始资金: {start_cash}\n回测期间：{start_date.strftime('%Y%m%d')}:{end_date.strftime('%Y%m%d')}")
print(f"总资金: {round(port_value, 2)}")
print(f"净收益: {round(pnl, 2)}")

# 绘制图表
cerebro.plot(style='candlestick')